﻿using System;
using System.Collections.Generic;
using Algoritmia.DataStructures.PriorityMaps;
using Sets = Algoritmia.DataStructures.Sets;
using Algoritmia.DataStructures.Digraphs;

namespace Algoritmia.Problems.SpanningTrees.Prim
{
    public class PrimsWithPriorityDictMinimumSpanningForestFinder<T> : PrimsMinimumSpanningForestFinder<T>
    {
        private readonly Func<IEnumerable<T>, IPriorityMap<T, Tuple<double, T>>> _createPriorityMap;

        public PrimsWithPriorityDictMinimumSpanningForestFinder(Func<IEnumerable<T>, Sets.ISet<T>> createSet = null,
                                                          Func<IEnumerable<T>,IPriorityMap<T,Tuple<double,T>>> createFringe = null)
            : base(createSet)
        {
            _createPriorityMap = createFringe ?? (it => new HeapMap<T, Tuple<double, T>>(comparison: (kv1, kv2) => kv1.Value.Item1.CompareTo(kv2.Value.Item1)));
        }

        protected override IEnumerable<Tuple<T,T>> MST(IDigraph<T> g, WeightingFunction<T,double> d, T u, Sets.ISet<T> added )
        {
            added.Add(u);
            var inFrontierFrom = _createPriorityMap(g.V);
            foreach (var v in g.Successors(u))
            {
                inFrontierFrom[v] = new Tuple<double, T>(d(u,v),u);
            }
            while (inFrontierFrom.Count > 0)
            {
                var optimalTuple = inFrontierFrom.ExtractOptimal();
                var v = optimalTuple.Key;
                u = optimalTuple.Value.Item2;
                var dEdge = optimalTuple.Value.Item1;
                added.Add(v);
                yield return new Tuple<T, T>(u, v);
                foreach (var w in g.Successors(v))
                {
                    if (!added.Contains(w) && (!inFrontierFrom.Contains(w) || d(inFrontierFrom[w].Item2, w) > d(v, w)))
                    {
                        inFrontierFrom[w] = new Tuple<double, T>(d(v, w),v);
                    }
                }
            }
        }
    }
}
